import os
import torch
import cv2 as cv
import numpy as np
from torch.utils.data.dataset import Dataset


class VideoDataset(Dataset):
    def __init__(self, frame_len=31):
        super().__init__()
        self.videos = []
        prepath = 'E:/dataset/ucf101'
        action = os.listdir(prepath)
        for label, path in enumerate(action):
            action_path = prepath + '/' + path
            video_name = os.listdir(action_path)
            for video in video_name:
                video_path = action_path + '/' + video
                self.videos.append(VideoInfo(video_path, label, frame_len))

    def __len__(self):
        total_len = 0
        for i in range(len(self.videos)):
            total_len += len(self.videos[i])
        return total_len

    def __getitem__(self, item):
        for i in range(len(self.videos)):
            if item < len(self.videos[i]):
                return self.videos[i][item]
            item -= len(self.videos[i])


class VideoInfo(Dataset):
    def __init__(self, video_path, label, frame_len):
        self.videopath = video_path
        self.video = cv.VideoCapture(video_path)
        assert self.video.isOpened() is True, 'video error.'
        self.cntframe = int(self.video.get(cv.CAP_PROP_FRAME_COUNT))
        self.label = torch.tensor(label, dtype=torch.long)
        self.frame_len = frame_len

    def __len__(self):
        cnt = self.cntframe - self.frame_len - 3
        assert cnt >= 1, 'Video too short.'
        return cnt

    def __getitem__(self, item):
        ret = self.video.set(cv.CAP_PROP_POS_FRAMES, item)
        assert ret is True, 'unknown3.'
        inputs = torch.zeros([self.frame_len, 240, 320, 3])
        cnt = 0
        while self.video.isOpened():
            ret, image_origin = self.video.read()
            assert ret is True, 'unknown1.'
            image = torch.tensor(image_origin / 256, dtype=torch.float32).unsqueeze(0)
            inputs[cnt] = image
            cnt += 1
            if cnt == self.frame_len:
                inputs = inputs.permute(3, 0, 1, 2)
                image_test(inputs)
                return inputs, self.label
        raise Exception('unkonwn2.')


def image_test(data):
    data = data.permute(1, 2, 3, 0)*256
    for n in range(len(data)):
        cv.imshow('test', data[n].numpy().astype(np.int8))
        cv.waitKey(10)
